import random
import torch


def synthetic_data(w, b, num):
    # generate y = Xw + b +noise
    X = torch.normal(0, 1, (num, len(w)))
    X = torch.as_tensor(X, dtype=torch.float32)
    w = torch.as_tensor(w, dtype=torch.float32)
    w = w.reshape((len(w), 1))
    Y = torch.matmul(X, w) + b
    Y += torch.normal(0, 0.01, Y.shape)
    return X, Y.reshape([num, 1])


# true_w = torch.tensor([2, -3, 4])
# true_b = 4.2
# features, labels = synthetic_data(true_w, true_b, 1000)
# print(features[0], labels[0])
# print(labels.shape)
# d2l.set_figsize()
# d2l.plt.scatter(features[:, 1].detach().numpy(), labels.detach().numpy(), 1)
# d2l.plt.show()


# load a batch of data
def data_iter(batch_size, feature, label):
    num_exam = len(feature)
    indices = list(range(num_exam))
    random.shuffle(indices)
    for i in range(0, num_exam, batch_size):
        batch_indices = torch.tensor(indices[i:min(i + batch_size, num_exam)])
        yield feature[batch_indices], label[batch_indices]


# define linear regression function
def linear_regression(X, w, b):
    result = torch.matmul(X, w) + b
    return result


# define loss function
def square_loss(y_hat, y):
    result = (y_hat - y.reshape(y_hat.shape)) ** 2 / 2
    return result


# Define optimization algorithm
def sgd(params, lr, batch_size):
    with torch.no_grad():
        for param in params:
            param -= lr * param.grad / batch_size
            param.grad.zero_()


# initialize model parameters
true_w = torch.tensor([2, -3, 4])
true_b = 4.2
features, labels = synthetic_data(true_w, true_b, 1000)
w = torch.normal(0, 0.01, size=(3, 1), requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.03
batch_size = 10
num_epochs = 1000
net = linear_regression
loss = square_loss
for epoch in range(num_epochs):
    for X, y in data_iter(batch_size, features, labels):
        l = loss(net(X, w, b), y)
        l.sum().backward()
        sgd([w, b], lr, batch_size)
    with torch.no_grad():
        train_l = loss(net(features, w, b), labels)
        print(f'epoch{epoch + 1},loss{float(train_l.mean()):f}')
